import os
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import numpy as np
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from torchvision.models import VGG16_Weights
import pickle
import time
from tkinter import ttk
import cv2
import concurrent.futures
import threading
from skimage.feature import local_binary_pattern


os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

print(">>>>>>>>>>>>>>>>>>start<<<<<<<<<<<<<<<<<<<")


class ImageMatcher(tk.Frame):
    def __init__(self, master=None, **kw):
        super().__init__(master, **kw)
        self.master = master
        self.start_x = None
        self.start_y = None
        self.rect = None
        self.selection_mode = False
        self.current_image = None
        self.features_cache = {}  # 特征缓存字典
        self.features_cache_file = "features_cache.pkl"  # 特征缓存文件

        # 添加特征索引
        self.feature_index = {}  # 用于存储预处理的特征
        self.feature_index_file = "feature_index.pkl"  # 特征索引文件

        self.create_widgets()
        self.init_vgg_model()
        self.load_features_cache()  # 加载特征缓存
        self.load_feature_index()  # 加载特征索引

        # 修改特征权重设置
        self.weights = {
            'vgg': 0.2,  # 降低VGG特征权重
            'stroke': 0.5,  # 提高笔画特征权重
            'stroke': 0.5,  # 提高笔画特征权重
            'structure': 0.3  # 保持结构特征权重
        }

    def load_feature_index(self):
        """加载特征索引"""
        if os.path.exists(self.feature_index_file):
            try:
                with open(self.feature_index_file, 'rb') as f:
                    self.feature_index = pickle.load(f)
                print(f"已加载特征索引，包含 {len(self.feature_index)} 个图像特征")
                print(f"特征索引中的文件列表：")
                for filename in self.feature_index.keys():
                    print(f"  - {filename}")
            except Exception as e:
                print(f"加载特征索引失败: {e}")
                self.feature_index = {}
        else:
            print("特征索引文件不存在，将创建新的索引")
            self.feature_index = {}

    def save_feature_index(self):
        """保存特征索引"""
        try:
            # 确保所有特征向量具有相同的维度
            feature_dim = None
            valid_features = {}
            
            for filename, features in self.feature_index.items():
                if features is not None:
                    if feature_dim is None:
                        feature_dim = features.shape[0]
                    if features.shape[0] == feature_dim:
                        valid_features[filename] = features
            
            self.feature_index = valid_features
            
            # 保存特征索引
            with open(self.feature_index_file, 'wb') as f:
                pickle.dump(self.feature_index, f)
            print(f"特征索引已保存，包含 {len(self.feature_index)} 个有效特征")
            
            # 保存特征缓存
            with open(self.features_cache_file, 'wb') as f:
                pickle.dump(self.features_cache, f)
            print(f"特征缓存已保存，包含 {len(self.features_cache)} 个特征")
            
        except Exception as e:
            print(f"保存特征索引失败: {e}")

    def init_vgg_model(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")
        self.model = models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1)
        self.model.classifier = self.model.classifier[:-1]
        self.model.eval()
        self.model = self.model.to(self.device)

        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])

        # 添加简单特征提取的变换
        self.simple_transform = transforms.Compose([
            transforms.Resize((64, 64)),
            transforms.ToTensor(),
        ])

    def create_widgets(self):
        # 创建左半部分控件
        self.left_frame = tk.Frame(self)
        self.left_frame.pack(side='left', fill='both', expand=True)

        self.image_frame = tk.Frame(self.left_frame, bg='gray')
        self.image_frame.pack(side='top', fill='both', expand=True, padx=10, pady=10)

        self.canvas = tk.Canvas(self.image_frame, bg='gray')
        self.canvas.pack(fill='both', expand=True)

        self.button_frame = tk.Frame(self.left_frame)
        self.button_frame.pack(side='left', padx=10, pady=10)

        self.upload_button = tk.Button(self.button_frame, text='上传图像', command=self.load_image)
        self.upload_button.pack(side='left', padx=5)

        self.match_button = tk.Button(self.button_frame, text='图像检索', command=self.match_images)
        self.match_button.pack(side='left', padx=5)

        self.local_match_button = tk.Button(self.button_frame, text='局部检索', command=self.start_local_selection)
        self.local_match_button.pack(side='left', padx=5)

        self.add_button = tk.Button(self.button_frame, text='添加图片', command=self.add_selection)
        self.add_button.pack(side='left', padx=5)

        self.remove_button = tk.Button(self.button_frame, text='删除图片', command=self.remove_selection)
        self.remove_button.pack(side='left', padx=5)

        # 添加预处理特征按钮
        self.preprocess_button = tk.Button(self.button_frame, text='预处理特征', command=self.preprocess_features)
        self.preprocess_button.pack(side='left', padx=5)

        # 创建右半部分控件
        self.right_frame = tk.Frame(self)
        self.right_frame.pack(side='left', fill='both', expand=True)

        # 创建5个图片框和相似度标签
        self.image_frames = []
        self.similarity_labels = []

        for i in range(5):
            frame = tk.Frame(self.right_frame, bg='gray')
            frame.pack(side='top', fill='both', expand=True, padx=10, pady=5)
            self.image_frames.append(frame)

            label = tk.Label(self.right_frame, text='相似度：0.00')
            label.pack(side='top', padx=10, pady=2)
            self.similarity_labels.append(label)

        # 添加状态栏
        self.status_bar = tk.Label(self, text="就绪", bd=1, relief=tk.SUNKEN, anchor=tk.W)
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)

    def load_features_cache(self):
        """加载特征缓存"""
        if os.path.exists(self.features_cache_file):
            try:
                with open(self.features_cache_file, 'rb') as f:
                    self.features_cache = pickle.load(f)
                self.status_bar.config(text=f"已加载特征缓存，包含 {len(self.features_cache)} 个图像特征")
                print(f"已加载特征缓存，包含 {len(self.features_cache)} 个图像特征")
            except Exception as e:
                print(f"加载特征缓存失败: {e}")
                self.features_cache = {}
        else:
            self.features_cache = {}

    def save_features_cache(self):
        """保存特征缓存"""
        try:
            with open(self.features_cache_file, 'wb') as f:
                pickle.dump(self.features_cache, f)
            print(f"特征缓存已保存，包含 {len(self.features_cache)} 个图像特征")
        except Exception as e:
            print(f"保存特征缓存失败: {e}")

    def extract_features(self, image_path):
        """提取图像特征，优先从缓存中获取"""
        # 检查缓存中是否已有该图像的特征
        if image_path in self.features_cache:
            return self.features_cache[image_path]

        try:
            image = Image.open(image_path).convert('RGB')
            image = self.transform(image).unsqueeze(0).to(self.device)

            with torch.no_grad():
                features = self.model(image)

            # 将特征保存到缓存
            features_np = features.cpu().numpy()
            self.features_cache[image_path] = features_np
            return features_np
        except Exception as e:
            print(f"特征提取错误: {e}")
            return None

    def extract_simple_features(self, image_path):
        """提取简单的像素特征，适用于局部检索"""
        try:
            image = Image.open(image_path).convert('L')  # 转为灰度图
            # 保持原始宽高比
            w, h = image.size
            # 统一大小为64x64
            image = image.resize((64, 64))
            # 转为numpy数组并归一化
            features = np.array(image).flatten() / 255.0
            return features
        except Exception as e:
            print(f"简单特征提取错误: {e}")
            return None

    def extract_stroke_features(self, image_path):
        """改进的汉字笔画特征提取方法"""
        try:
            # 修正文件路径
            image_path = os.path.normpath(image_path)

            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                try:
                    pil_img = Image.open(image_path).convert('RGB')
                    img = np.array(pil_img)
                    img = img[:, :, ::-1].copy()
                except Exception as e:
                    print(f"无法读取图像: {e}")
                    return None

            # 1. 改进的图像预处理
            img = cv2.resize(img, (128, 128))  # 保持较大的尺寸以保留更多细节
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            
            # 使用自适应阈值进行二值化，更好地处理不同光照条件
            binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                        cv2.THRESH_BINARY_INV, 11, 2)
            
            # 降噪处理
            kernel = np.ones((2, 2), np.uint8)
            binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
            binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)

            # 2. 提取骨架
            try:
                skeleton = cv2.ximgproc.thinning(binary)
            except:
                # 如果没有ximgproc，使用形态学操作
                kernel = np.ones((3, 3), np.uint8)
                skeleton = binary.copy()
                while True:
                    eroded = cv2.erode(skeleton, kernel)
                    temp = cv2.dilate(eroded, kernel)
                    temp = cv2.subtract(skeleton, temp)
                    skeleton = eroded.copy()
                    if cv2.countNonZero(temp) == 0:
                        break

            # 3. 提取改进的特征
            # 3.1 计算笔画方向特征（使用16个方向而不是8个，提高精度）
            gx = cv2.Sobel(skeleton, cv2.CV_32F, 1, 0, ksize=3)
            gy = cv2.Sobel(skeleton, cv2.CV_32F, 0, 1, ksize=3)
            mag, ang = cv2.cartToPolar(gx, gy)
            
            bins = 16  # 增加方向数量
            angle_bins = np.zeros(bins)
            for i in range(skeleton.shape[0]):
                for j in range(skeleton.shape[1]):
                    if skeleton[i, j] > 0:
                        bin_idx = int(ang[i, j] * bins / (2 * np.pi))
                        if bin_idx >= bins:
                            bin_idx = bins - 1
                        angle_bins[bin_idx] += 1

            # 归一化方向特征
            if np.sum(angle_bins) > 0:
                angle_bins = angle_bins / np.sum(angle_bins)

            # 3.2 提取改进的区域特征（使用4x4网格而不是3x3，提供更细致的空间信息）
            grid_size = 4
            grid_features = np.zeros(grid_size * grid_size)
            h, w = binary.shape
            cell_h, cell_w = h // grid_size, w // grid_size

            for i in range(grid_size):
                for j in range(grid_size):
                    cell = binary[i*cell_h:(i+1)*cell_h, j*cell_w:(j+1)*cell_w]
                    grid_features[i * grid_size + j] = np.sum(cell > 0) / (cell_h * cell_w)

            # 3.3 提取笔画交叉点和端点特征
            kernel = np.ones((3, 3), np.uint8)
            dilated = cv2.dilate(skeleton, kernel, iterations=1)
            endpoints = cv2.subtract(skeleton, cv2.erode(skeleton, kernel))
            crosspoints = cv2.bitwise_and(dilated, skeleton)
            
            num_endpoints = np.sum(endpoints > 0) / 255
            num_crosspoints = np.sum(crosspoints > 0) / 255

            # 3.4 计算笔画复杂度特征
            contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if contours:
                max_contour = max(contours, key=cv2.contourArea)
                perimeter = cv2.arcLength(max_contour, True)
                area = cv2.contourArea(max_contour)
                compactness = perimeter * perimeter / (4 * np.pi * area) if area > 0 else 0
                
                # 计算轮廓的Hu矩
                moments = cv2.moments(max_contour)
                hu_moments = cv2.HuMoments(moments).flatten()[:4]  # 使用前4个Hu矩
            else:
                compactness = 0
                hu_moments = np.zeros(4)

            # 4. 组合所有特征
            features = np.concatenate([
                [num_endpoints, num_crosspoints, compactness],  # 基本特征
                angle_bins,                                     # 方向特征（16维）
                grid_features,                                  # 区域特征（16维）
                hu_moments,                                     # 形状特征（4维）
            ])

            # 5. 特征归一化
            features = (features - np.min(features)) / (np.max(features) - np.min(features) + 1e-10)
            
            return features

        except Exception as e:
            print(f"特征提取错误: {e}")
            return None

    def extract_enhanced_features(self, image_path):
        """增强的特征提取方法，确保所有特征维度一致"""
        try:
            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                try:
                    pil_img = Image.open(image_path).convert('RGB')
                    img = np.array(pil_img)
                    img = img[:, :, ::-1].copy()
                except Exception as e:
                    print(f"无法读取图像: {e}")
                    return None
            
            # 1. 图像预处理
            img = cv2.resize(img, (64, 64))  # 统一大小
            img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img_binary = cv2.adaptiveThreshold(img_gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                             cv2.THRESH_BINARY_INV, 11, 2)
            
            # 2. 提取HOG特征
            hog = cv2.HOGDescriptor((64,64), (16,16), (8,8), (8,8), 9)
            hog_features = hog.compute(img_gray).flatten()
            
            # 3. 提取LBP特征
            radius = 1
            n_points = 8 * radius
            lbp = local_binary_pattern(img_gray, n_points, radius, method='uniform')
            lbp_hist = np.histogram(lbp.ravel(), bins=np.arange(0, n_points + 3), 
                                   range=(0, n_points + 2), density=True)[0]
            
            # 4. 提取轮廓和Hu矩特征
            contours, _ = cv2.findContours(img_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if contours:
                max_contour = max(contours, key=cv2.contourArea)
                moments = cv2.moments(max_contour)
                hu_moments = cv2.HuMoments(moments).flatten()
            else:
                hu_moments = np.zeros(7)
            
            # 5. 提取方向特征
            gx = cv2.Sobel(img_binary, cv2.CV_32F, 1, 0, ksize=3)
            gy = cv2.Sobel(img_binary, cv2.CV_32F, 0, 1, ksize=3)
            mag, ang = cv2.cartToPolar(gx, gy)
            direction_hist = np.histogram(ang, bins=16, density=True)[0]
            
            # 6. 组合所有特征
            combined_features = np.concatenate([
                hog_features,      # HOG特征
                lbp_hist,         # LBP特征
                hu_moments,       # Hu矩特征
                direction_hist    # 方向特征
            ])
            
            # 7. 标准化特征向量
            if np.any(combined_features):
                combined_features = combined_features / (np.linalg.norm(combined_features) + 1e-10)
            
            return combined_features
        
        except Exception as e:
            print(f"特征提取错误: {e}")
            return None

    def preprocess_features(self):
        """使用多线程预处理所有图像特征"""
        # 清理旧的特征文件
        if os.path.exists(self.feature_index_file):
            os.remove(self.feature_index_file)
        if os.path.exists(self.features_cache_file):
            os.remove(self.features_cache_file)
        
        folder = r"image_label"

        if not os.path.exists(folder):
            messagebox.showwarning("警告", f"找不到 {folder} 文件夹")
            return

        filenames = []
        for root, dirs, files in os.walk(folder):
            for file in files:
                if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    filenames.append(os.path.join(root, file))

        if not filenames:
            messagebox.showinfo("信息", "文件夹中没有找到图像文件")
            return

        print(f"开始预处理 {len(filenames)} 个图像")
        self.status_bar.config(text=f"正在预处理 {len(filenames)} 个图像特征...")

        progress_window = tk.Toplevel(self.master)
        progress_window.title("预处理进度")
        progress_window.geometry("400x150")
        progress_window.resizable(False, False)

        self.cancel_preprocessing = False

        def cancel_process():
            self.cancel_preprocessing = True
            cancel_button.config(state='disabled')
            progress_label.config(text="正在取消...")

        cancel_button = tk.Button(progress_window, text="取消", command=cancel_process)
        cancel_button.pack(pady=5)

        progress_var = tk.IntVar()
        progress_label = tk.Label(progress_window, text="0%")
        progress_label.pack(pady=5)

        progress_bar = ttk.Progressbar(progress_window, variable=progress_var,
                                       maximum=len(filenames))
        progress_bar.pack(pady=5, padx=10, fill='x')

        info_label = tk.Label(progress_window, text="正在处理...", wraplength=380)
        info_label.pack(pady=5)

        start_time = time.time()
        processed_count = 0
        success_count = 0

        def process_image(filename):
            """处理单个图像"""
            try:
                if filename not in self.feature_index:
                    # 使用增强特征提取方法替代原来的extract_features
                    features = self.extract_enhanced_features(filename)
                    if features is not None:
                        return filename, features, True
                return filename, self.feature_index.get(filename), True
            except Exception as e:
                print(f"处理图片出错 {filename}: {e}")
                return filename, None, False

        def update_progress(filename, success):
            """更新进度显示"""
            nonlocal processed_count, success_count
            processed_count += 1
            if success:
                success_count += 1

            progress_var.set(processed_count)
            percent = (processed_count / len(filenames)) * 100
            elapsed_time = time.time() - start_time
            avg_time = elapsed_time / processed_count if processed_count > 0 else 0
            eta = avg_time * (len(filenames) - processed_count)

            status_text = (f"进度: {processed_count}/{len(filenames)} ({percent:.1f}%)\n"
                           f"成功: {success_count} 个\n"
                           f"预计剩余时间: {eta:.1f} 秒")
            progress_label.config(text=status_text)

            base_filename = os.path.basename(filename)
            info_text = f"正在处理: {base_filename}"
            info_label.config(text=info_text)

            progress_window.update()

        # 使用线程池处理图像
        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
            future_to_filename = {
                executor.submit(process_image, filename): filename
                for filename in filenames
            }

            for future in concurrent.futures.as_completed(future_to_filename):
                if self.cancel_preprocessing:
                    for f in future_to_filename:
                        f.cancel()
                    break

                filename = future_to_filename[future]
                try:
                    filename, features, success = future.result()
                    if features is not None:
                        self.feature_index[filename] = features
                    update_progress(filename, success)
                except Exception as e:
                    print(f"处理文件 {filename} 时出错: {e}")
                    update_progress(filename, False)

        # 保存特征索引
        if not self.cancel_preprocessing:
            self.save_feature_index()
            final_message = (f"预处理完成\n"
                             f"总文件数: {len(filenames)}\n"
                             f"成功处理: {success_count}\n"
                             f"总耗时: {time.time() - start_time:.1f} 秒")
        else:
            final_message = "预处理已取消"

        progress_window.destroy()

        if not self.cancel_preprocessing:
            print(f"预处理完成，特征索引包含 {len(self.feature_index)} 个图像特征")
            messagebox.showinfo("完成", final_message)
        else:
            print("预处理已取消")
            messagebox.showinfo("已取消", final_message)

        self.status_bar.config(text=final_message.replace('\n', ' | '))

    def calculate_enhanced_similarity(self, features1, features2):
        """计算两个特征向量之间的相似度"""
        try:
            if features1 is None or features2 is None:
                return 0.0
            
            # 确保特征向量形状一致
            features1 = features1.flatten()
            features2 = features2.flatten()
            
            if features1.shape != features2.shape:
                print(f"特征维度不匹配: {features1.shape} vs {features2.shape}")
                return 0.0
            
            # 使用余弦相似度
            similarity = np.dot(features1, features2) / (
                np.linalg.norm(features1) * np.linalg.norm(features2) + 1e-10
            )
            
            return float(similarity)
        
        except Exception as e:
            print(f"相似度计算错误: {e}")
            return 0.0

    def add_selection(self):
        file_path = filedialog.askopenfilename(
            title="选择图片",
            filetypes=[("图片文件", "*.jpg;*.jpeg;*.png;*.bmp;*.gif")]
        )
        if not file_path:
            print("未选择任何文件")
            return

        try:
            image = Image.open(file_path)
            print(f"成功加载图片：{file_path}")

            save_path = filedialog.asksaveasfilename(
                defaultextension=".png",
                filetypes=[("PNG文件", "*.png"), ("JPEG文件", "*.jpg"), ("所有文件", "*.*")]
            )

            if save_path:
                image.save(save_path)
                # 提取并缓存新添加图片的特征
                self.extract_features(save_path)
                self.save_features_cache()
                print(f"图片已成功保存到：{save_path}")
                messagebox.showinfo("成功", "图片已成功添加到数据库")
            else:
                print("未选择保存路径")
        except Exception as e:
            print(f"处理图片失败：{e}")
            messagebox.showerror("错误", f"处理图片失败：{e}")

    def remove_selection(self):
        file_path = filedialog.askopenfilename(
            title="选择要删除的图片",
            filetypes=[("图片文件", "*.jpg;*.jpeg;*.png;*.bmp;*.gif")]
        )
        if not file_path:
            print("未选择任何文件")
            return

        confirm = messagebox.askyesno("确认删除", f"确定要删除文件：\n{file_path} 吗？")
        if confirm:
            try:
                os.remove(file_path)
                # 从缓存中删除该图片的特征
                if file_path in self.features_cache:
                    del self.features_cache[file_path]
                    self.save_features_cache()
                print(f"文件已成功删除：{file_path}")
                messagebox.showinfo("删除成功", f"文件已成功删除：\n{file_path}")
            except Exception as e:
                print(f"删除文件失败：{e}")
                messagebox.showerror("删除失败", f"删除文件失败：\n{e}")
        else:
            print("删除操作已取消")

    def load_image(self):
        filetypes = (("JPEG files", "*.jpg"), ("PNG files", "*.png"))
        self.imgName = filedialog.askopenfilename(title="选择图像文件", filetypes=filetypes)

        if self.imgName:
            self.clear_all_images()  # 使用新方法清除所有图像
            image = Image.open(self.imgName)
            self.original_image = image
            image = self.resize_image(image, self.canvas.winfo_width(), self.canvas.winfo_height())
            self.current_image = ImageTk.PhotoImage(image)
            self.canvas.create_image(0, 0, anchor='nw', image=self.current_image)

    def clear_images(self):
        # 只清除右侧的结果图像，不清除左侧的查询图像
        for label in self.similarity_labels:
            label.config(text='相似度：0.00')

        for frame in self.image_frames:
            for widget in frame.winfo_children():
                widget.destroy()

    def clear_all_images(self):
        # 清除所有图像，包括左侧的查询图像
        self.canvas.delete("all")
        self.clear_images()

    def resize_image(self, image, width, height):
        if width <= 1 or height <= 1:  # 防止除零错误
            return image
        w, h = image.size
        ratio = min(width / w, height / h)
        size = (int(w * ratio), int(h * ratio))
        return image.resize(size)

    def start_local_selection(self):
        self.selection_mode = True
        self.canvas.bind('<Button-1>', self.on_mouse_down)
        self.canvas.bind('<B1-Motion>', self.on_mouse_drag)
        self.canvas.bind('<ButtonRelease-1>', self.on_mouse_release)

    def on_mouse_down(self, event):
        self.start_x = event.x
        self.start_y = event.y
        if self.rect:
            self.canvas.delete(self.rect)

    def on_mouse_drag(self, event):
        if self.rect:
            self.canvas.delete(self.rect)
        self.rect = self.canvas.create_rectangle(
            self.start_x, self.start_y, event.x, event.y,
            outline='red', width=2
        )

    def on_mouse_release(self, event):
        if self.rect:
            x1 = min(self.start_x, event.x)
            y1 = min(self.start_y, event.y)
            x2 = max(self.start_x, event.x)
            y2 = max(self.start_y, event.y)

            # 裁剪选择区域
            region = self.original_image.crop((x1, y1, x2, y2))
            temp_path = "temp_region.png"

            # 保持原始宽高比，添加白色背景填充而不是拉伸
            w, h = region.size
            max_size = max(w, h)
            padded_region = Image.new('RGB', (max_size, max_size), (255, 255, 255))
            padded_region.paste(region, ((max_size - w) // 2, (max_size - h) // 2))
            padded_region.save(temp_path)

            # 显示选择的区域信息
            print(f"选择区域大小: {w}x{h}")

            # 使用改进的局部检索方法
            self.match_local_images_stroke_based(temp_path)

            os.remove(temp_path)

        self.selection_mode = False
        self.canvas.unbind('<Button-1>')
        self.canvas.unbind('<B1-Motion>')
        self.canvas.unbind('<ButtonRelease-1>')

    def match_local_images_stroke_based(self, local_image_path):
        """改进的局部检索方法，使用增强特征提取"""
        folder = r"image_label"

        if not os.path.exists(folder):
            messagebox.showwarning("警告", f"找不到 {folder} 文件夹，请确保该文件夹存在")
            return

        if len(self.feature_index) == 0:
            messagebox.showwarning("警告", "特征索引为空，请先点击'预处理特征'按钮")
            return

        print(f"开始局部检索...")
        start_time = time.time()
        self.status_bar.config(text="正在提取查询图片特征...")
        self.update()

        # 提取查询图像的特征
        query_features = self.extract_enhanced_features(local_image_path)
        if query_features is None:
            messagebox.showerror("错误", "无法提取查询图片的特征")
            return
        
        # 验证特征维度
        feature_dim = query_features.shape[0]
        valid_features = {k: v for k, v in self.feature_index.items() 
                         if v is not None and v.shape[0] == feature_dim}
        
        if not valid_features:
            messagebox.showwarning("警告", "没有找到维度匹配的特征，请重新预处理特征")
            return

        self.status_bar.config(text="正在计算相似度...")
        self.update()

        # 创建进度窗口
        progress_window = tk.Toplevel(self.master)
        progress_window.title("局部检索进度")
        progress_window.geometry("400x100")

        progress_var = tk.IntVar()
        progress_label = tk.Label(progress_window, text="计算相似度...")
        progress_label.pack(pady=10)
        progress_bar = ttk.Progressbar(progress_window, variable=progress_var,
                                       maximum=len(valid_features))
        progress_bar.pack(pady=10)

        similarities = []
        
        def calculate_similarity_task(filename, features):
            """单个相似度计算任务"""
            try:
                # 计算相似度
                similarity = self.calculate_enhanced_similarity(query_features, features)
                
                # 额外计算结构相似度
                img1 = cv2.imread(local_image_path, cv2.IMREAD_GRAYSCALE)
                img2 = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
                
                if img1 is not None and img2 is not None:
                    # 调整大小
                    img1 = cv2.resize(img1, (64, 64))
                    img2 = cv2.resize(img2, (64, 64))
                    
                    # 二值化
                    _, img1_bin = cv2.threshold(img1, 127, 255, cv2.THRESH_BINARY)
                    _, img2_bin = cv2.threshold(img2, 127, 255, cv2.THRESH_BINARY)
                    
                    # 计算结构相似度（使用像素重叠度）
                    intersection = cv2.bitwise_and(img1_bin, img2_bin)
                    union = cv2.bitwise_or(img1_bin, img2_bin)
                    if cv2.countNonZero(union) > 0:
                        structural_similarity = cv2.countNonZero(intersection) / cv2.countNonZero(union)
                    else:
                        structural_similarity = 0
                    
                    # 组合两种相似度
                    combined_similarity = 0.4 * similarity + 0.6 * structural_similarity
                    return filename, combined_similarity
                
                return filename, similarity
                
            except Exception as e:
                print(f"计算相似度出错 {filename}: {e}")
                return None

        # 使用线程池进行并行计算
        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
            future_to_filename = {
                executor.submit(calculate_similarity_task, filename, features): filename
                for filename, features in valid_features.items()
            }

            for future in concurrent.futures.as_completed(future_to_filename):
                result = future.result()
                if result:
                    filename, similarity = result
                    # 降低相似度阈值，增加匹配结果
                    if similarity > 0.3:  # 降低阈值到0.3
                        similarities.append((filename, similarity))
                progress_var.set(len(similarities))
                progress_window.update()

        progress_window.destroy()

        if not similarities:
            messagebox.showinfo("提示", "未找到任何相似图片")
            self.status_bar.config(text="未找到任何相似图片")
            return

        # 按相似度排序并显示结果
        similarities.sort(key=lambda x: x[1], reverse=True)
        top_similarities = similarities[:5]  # 显示前5个最相似的结果
        
        self.clear_images()
        
        for i, (filename, similarity) in enumerate(top_similarities):
            try:
                image = Image.open(filename)
                image = self.resize_image(image,
                                          self.image_frames[i].winfo_width(),
                                          self.image_frames[i].winfo_height())
                photo = ImageTk.PhotoImage(image)
                label = tk.Label(self.image_frames[i], image=photo)
                label.image = photo
                label.pack(fill='both', expand=True)

                # 获取文件夹名和文件名
                folder_name = os.path.basename(os.path.dirname(filename))
                base_filename = os.path.basename(filename)
                name_without_ext = os.path.splitext(base_filename)[0]

                if len(name_without_ext) > 15:
                    display_name = name_without_ext[:12] + "..."
                else:
                    display_name = name_without_ext

                # 显示详细的相似度信息
                self.similarity_labels[i].config(
                    text=f'局部相似度：{similarity:.4f} | 文件夹：{folder_name} | 图片：{display_name}'
                )
            except Exception as e:
                print(f"显示图片 {filename} 时出错: {e}")

        elapsed_time = time.time() - start_time
        self.status_bar.config(
            text=f"局部检索完成，耗时 {elapsed_time:.2f} 秒"
        )

    def match_images(self, local_image_path=None, is_local=False):
        """改进的图像匹配方法"""
        if not hasattr(self, 'imgName') and not is_local:
            messagebox.showwarning("警告", "请先上传图像")
            return

        if len(self.feature_index) == 0:
            messagebox.showwarning("警告", "特征索引为空，请先点击'预处理特征'按钮")
            return

        start_time = time.time()
        self.status_bar.config(text="正在提取查询图片特征...")
        self.update()

        # 提取查询图片的特征
        source_image = local_image_path if is_local else self.imgName
        query_features = self.extract_enhanced_features(source_image)
        
        if query_features is None:
            messagebox.showerror("错误", "无法提取查询图片的特征")
            return
        
        # 验证特征维度
        feature_dim = query_features.shape[0]
        valid_features = {k: v for k, v in self.feature_index.items() 
                         if v is not None and v.shape[0] == feature_dim}
        
        if not valid_features:
            messagebox.showwarning("警告", "没有找到维度匹配的特征，请重新预处理特征")
            return

        self.status_bar.config(text="正在计算相似度...")
        self.update()

        # 创建进度窗口
        progress_window = tk.Toplevel(self.master)
        progress_window.title("检索进度")
        progress_window.geometry("400x100")

        progress_var = tk.IntVar()
        progress_label = tk.Label(progress_window, text="计算相似度...")
        progress_label.pack(pady=10)
        progress_bar = ttk.Progressbar(progress_window, variable=progress_var,
                                       maximum=len(self.feature_index))
        progress_bar.pack(pady=10)

        similarities = []
        
        def calculate_similarity_task(filename, features):
            """单个相似度计算任务"""
            try:
                similarity = self.calculate_enhanced_similarity(query_features, features)
                return filename, similarity
            except Exception as e:
                print(f"计算相似度出错 {filename}: {e}")
                return None

        # 使用线程池进行并行计算
        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
            future_to_filename = {
                executor.submit(calculate_similarity_task, filename, features): filename
                for filename, features in self.feature_index.items()
            }

            for future in concurrent.futures.as_completed(future_to_filename):
                result = future.result()
                if result:
                    filename, similarity = result
                    similarities.append((filename, similarity))
                progress_var.set(len(similarities))
                progress_window.update()

        progress_window.destroy()

        if not similarities:
            messagebox.showinfo("提示", "未找到任何相似图片")
            self.status_bar.config(text="未找到任何相似图片")
            return

        # 按相似度排序并显示结果
        similarities.sort(key=lambda x: x[1], reverse=True)
        top_similarities = similarities[:5]
        
        self.clear_images()
        
        for i, (filename, similarity) in enumerate(top_similarities):
            try:
                image = Image.open(filename)
                image = self.resize_image(image,
                                          self.image_frames[i].winfo_width(),
                                          self.image_frames[i].winfo_height())
                photo = ImageTk.PhotoImage(image)
                label = tk.Label(self.image_frames[i], image=photo)
                label.image = photo
                label.pack(fill='both', expand=True)

                # 获取文件夹名和文件名
                folder_name = os.path.basename(os.path.dirname(filename))
                base_filename = os.path.basename(filename)
                name_without_ext = os.path.splitext(base_filename)[0]

                if len(name_without_ext) > 15:
                    display_name = name_without_ext[:12] + "..."
                else:
                    display_name = name_without_ext

                # 显示详细的相似度信息
                self.similarity_labels[i].config(
                    text=f'总相似度：{similarity:.4f} | 文件夹：{folder_name} | 图片：{display_name}'
                )
            except Exception as e:
                print(f"显示图片 {filename} 时出错: {e}")

        elapsed_time = time.time() - start_time
        self.status_bar.config(
            text=f"检索完成，耗时 {elapsed_time:.2f} 秒，找到 {len(similarities)} 张相似图片"
        )


if __name__ == '__main__':
    root = tk.Tk()
    root.geometry('1200x800+400+200')  # 调整窗口大小以适应5张图片
    root.title('古籍汉字图像检索系统')

    app = ImageMatcher(master=root)
    app.pack(fill='both', expand=True)

    root.mainloop()